import os
import time 
import multiprocessing
from multiprocessing import Pool


# 给基因簇中的基因添加信息
def run_add_gene_info(cluster_table, gene_info_file, output_file):
    cmd = (
        "python ./02.add_gene_info.py -cluster " + cluster_table + 
        " -gene " + gene_info_file +
        " -output " + output_file
    )
    os.system(cmd)
    
    
if __name__ == "__main__":
    start_time = time.time()
    # 生成任务列表
    files = os.listdir("./01.cluster_result")
    path = [os.path.join("./01.cluster_result",x) for x in files]
    speces = [x.split("___")[0] for x in files]
    gene_info_files = []
    for i in speces:
        gene_info_path = "../imp_data/" + i + "/eggnog.emapper.annotations.all.txt"
        gene_info_files.append(gene_info_path)
    output_path = [os.path.join("./02.cluster_add_gene_info/", x + "___add_gene_info.xlsx") for x in speces]
    os.makedirs("./02.cluster_add_gene_info", exist_ok=True)
    task_list = list(zip(path, gene_info_files, output_path))
    # 多进程执行
    cpu_count = multiprocessing.cpu_count() - 1
    with Pool(100) as pool:
        pool.starmap(run_add_gene_info, task_list)
    end_time = time.time()
    total_time = end_time - start_time
    print(f"\n任务数量是 {len(task_list)} :\n", f"\t用时 {total_time}s\n")
